#include <arrow/api.h>
#include <arrow/io/api.h>
#include <parquet/arrow/writer.h>
#include <iostream>

int main() {
    // 1️⃣ 创建内存池
    arrow::MemoryPool* pool = arrow::default_memory_pool();

    // 2️⃣ 构造 Arrow 列
    std::shared_ptr<arrow::Array> id_array;
    std::shared_ptr<arrow::Array> name_array;

    arrow::Status st;

    arrow::Int64Builder id_builder(pool);
    st = id_builder.Append(1);      if (!st.ok()) { std::cerr << st.ToString(); return -1; }
    st = id_builder.Append(2);      if (!st.ok()) { std::cerr << st.ToString(); return -1; }
    st = id_builder.Append(3);      if (!st.ok()) { std::cerr << st.ToString(); return -1; }
    st = id_builder.Finish(&id_array); if (!st.ok()) { std::cerr << st.ToString(); return -1; }

    arrow::StringBuilder name_builder(pool);
    st = name_builder.Append("Alice");   if (!st.ok()) { std::cerr << st.ToString(); return -1; }
    st = name_builder.Append("Bob");     if (!st.ok()) { std::cerr << st.ToString(); return -1; }
    st = name_builder.Append("Charlie"); if (!st.ok()) { std::cerr << st.ToString(); return -1; }
    st = name_builder.Finish(&name_array); if (!st.ok()) { std::cerr << st.ToString(); return -1; }

    // 3️⃣ 构造 Schema
    auto schema = arrow::schema({
        arrow::field("id", arrow::int64()),
        arrow::field("name", arrow::utf8())
    });

    // 4️⃣ 构造 Table
    auto table = arrow::Table::Make(schema, {id_array, name_array});

    // 5️⃣ 打开输出文件（新版 API ✅）
    auto result = arrow::io::FileOutputStream::Open("example.parquet", /*append=*/false);
    if (!result.ok()) {
        std::cerr << "❌ 无法打开输出文件: " << result.status().ToString() << std::endl;
        return -1;
    }

    std::shared_ptr<arrow::io::FileOutputStream> outfile = *result;

    // 6️⃣ 写入 Parquet 文件
    parquet::WriterProperties::Builder builder;
    builder.compression(parquet::Compression::SNAPPY);
    std::shared_ptr<parquet::WriterProperties> properties = builder.build();

    auto status = parquet::arrow::WriteTable(*table, pool, outfile, 1024, properties);
    if (!status.ok()) {
        std::cerr << "❌ 写入 Parquet 失败: " << status.ToString() << std::endl;
        return -1;
    }

    std::cout << "✅ Parquet 文件写入成功: example.parquet" << std::endl;
    return 0;
}
